function [cmeans, labels, dists, info] = slkmeans(X, varargin)
%SLKMEANS Performs (Generalized) K-means clustering
%
% [ Syntax ]
%   - [cmeans, labels] = slkmeans(X, ...)
%   - [cmeans, labels, dists] = slkmeans(X, ...)
%   - [cmeans, labels, dists, info] = slkmeans(X, ...)
%
% [ Arguments ]
%   - X:            the sample matrix
%   - cmeans:       the mean vectors of clusters
%   - labels:       the clustered labels
%   - dists:        the distance metric values from the samples to the
%                   corresponding centroids.
%   - info:         the information of the iterative procedure,
%                   it is a struct with the following fields:
%                   \{:
%                       - nsamples:    the number of samples
%                       - ncluster:    the number of output clusters
%                       - initK:       the initial K value
%                       - niters:      the number of iterations
%                       - converged:   whether the procedure converged
%                       - changes:     a column vector containing the
%                                      number of changes in all iterations.
%                       - avgdists:    a column vector containing the
%                                      average distances from the samples
%                                      to the corresponding centroids
%                                      in all iterations.
%                   \:}
%                   You may use info.changes(end) to see the number of
%                   samples that change in the final iteration. And, you
%                   may use info.avgdists(end) to see the average distances
%                   in the output, which is the final objective value.
%   
% [ Description ]
%   - [cmeans, labels] = slkmeans(X, ...) performs K-means clustering on
%     the sample set given by X.
%
%     Suppose there are n samples in a d-dimensional subspace, then X
%     should be a d x n matrix. The number of clusters K is specified 
%     through the option K as introduced below.
%
%     In the output, cmeans would be a d x nc matrix (nc <= K) containing 
%     the centroid vectors of the clusters. When there are no clusters
%     being dropped during the procedure, nc == K in the output.
%     The labels will be a 1 x n row vector of the labels assigned to the
%     samples.
%
%     You can specify the following options to control the clustering.
%     \{:
%        - K:           the (maximum) number of clusters (default = 3).
%                   
%                       Please note that the option name is 'K', which is
%                       in upper-case. Don't use 'k' as the option name.
%
%        - init_cmeans: the initial cluster means (default = []).
%
%                       You may set this option to a d x K matrix with each
%                       column being an initial mean vector for a cluster.
%
%        - init_labels: the initial labels of the samples. (default = []).
%
%                       You may set this option to a vector of length n.
%                       The values of labels should be integers ranging
%                       from 1 to K.
%
%                       If the initial_labels is specified 
%        - weights:     the weights of the samples. (default = []).
%
%                       You may set this option to a vector of length n
%                       containing the weights of the samples. 
%
%                       The weights will be utilized when computing the
%                       means.
%
%        - handle_empty: the strategy to deal with the cluster which no
%                        samples is assigned to at an iteration.
%                        \{:
%                           - error:      raise an error message
%                           - retain:     keep its mean intact 
%                           - singleton:  create a new cluster with one
%                                         sample that is furthest from the 
%                                         corresponding mean
%                           - drop:       simply dropping that cluster
%                        \:}
%                       By default, it is set to 'singleton'.
%
%        - start:       the strategy to initialize the procedure
%                       \{:
%                           - rndsample:  randomly samples K vectors from
%                                         the input sample set.
%                           - farsample:  first selects the sample that is
%                                         furthest to all other samples,
%                                         and then iteratively selects the 
%                                         one that is furthest to all
%                                         selected samples.
%                       \:}
%                       default = 'rndsample'.
%
%        - metric:      the metric to evaluate the distance. The selection
%                       of the metric affects not only the assignment of
%                       sample labels but also the way to compute the
%                       centroid.
%                       \{:
%                           - sqdist:   squared Euclidean distance
%                                       
%                                       The centroid is the conventional
%                                       mean vector.
%
%                           - corrdist: the correlation distance, 
%                                       which equals 1 - normalized
%                                       correlation. (cosine of angle)
%
%                                       The centroid is the normalized mean 
%                                       of normalized vectors.
%
%                           - cityblk:  the city block distance.
%
%                                       The centroid is a vector in which
%                                       each element is the median of the
%                                       corresponding elements of the
%                                       vectors in the cluster.
%                       \:}
%                       Default = 'sqdist'.
%
%        - compute:     You can also specify this option in form of 
%                       {f_metric, f_centroid} to give your own way to
%                       computing. This option facilitates the
%                       implementation of generalized k-means.
%
%                       Here, both f_metric and f_centroid are function
%                       handle.
%                          - f_metric should support the following syntax
%                               $ D = f_metric(X1, X2) $
%                            to compute pairwise distance metrics.
%                            The input X1 and X2 will be d x n1 and d x
%                            n2 sample matrices, while the output should be
%                            an n1 x n2 matrix of metric values.
%
%                          - f_centroid should support the following syntax
%                                $ v = f_centroid(X, w) $
%                            to compute the centroid of the samples in X
%                            (when w is empty), or the weighted centroid of
%                            the samples in X (when w is a row vector of
%                            weights).
%
%                       By default, it is set to [], which means to use
%                       standard way of computing.
%
%        - maxiter:     The maximum number of iterations. default = 100.
%
%        - verbose:     whether to show information in the procedure.
%                       default = true.
%     \:}
%
%   - [cmeans, labels, dists] = slkmeans(X, ...) additionally returns the 
%     distances from the samples to the corresponding centroid vectors.
%     The distance computation is specified by the option metric, or the
%     f_centroid function handle in the option compute. 
%   
%     By default, the distances are squared Euclidean distances.
%
%   - [cmeans, labels, dists, info] = slkmeans(X, ...) additionally returns 
%     the information of the iteration procedure.
%
% [ Remarks ]
%   - The function performs the clustering by iteratively updating the
%     centroids and the label assignment until convergence of maximum
%     iteration is reached.
%
%   - If either init_cmeans or init_labels is specified, it will be used
%     to initialize the clustering. Please note that only one of them can
%     be specified as non-empty. If both of them are not given, the
%     function initializes the procedure according to the option start, by
%     default, it randomly samples K vectors as initial centroids.
%
%   - The label assignment and centroid computation is decided by the
%     option metric, which by default is 'sqdist'. If the option compute is
%     specified, the option metric will be ignored, and the computation
%     is based on the function handles given in the option compute.
%
% [ History ]
%   - Created by Dahua Lin, on Jul 24, 2007
%

%% parse and verify input arguments

error(nargchk(1, inf, nargin));

% X
assert(isnumeric(X) && ndims(X) == 2, 'sltoolbox:slkmeans:invalidarg', ...
    'X should be a 2D numeric matrix.');
[d, n] = size(X);

% options

opts = struct( ...
    'K', 3, ...
    'init_cmeans', [], ...
    'init_labels', [], ...
    'weights', [], ...
    'handle_empty', 'singleton', ...
    'start', 'rndsample', ...
    'metric', 'sqdist', ...
    'compute', [], ...
    'maxiter', 100, ...
    'verbose', true);

if ~isempty(varargin)
    opts = setopts(opts, varargin{:});
end

% K
K = opts.K;
assert(isnumeric(K) && isscalar(K) && K > 0 && K == fix(K), ...
    'sltoolbox:slkmeans:invalidopt', 'K should be a positive integer scalar.');
assert(K <= n, ...
    'sltoolbox:slkmeans:invalidopt', 'K exceeds the number of samples.');

% init_cmeans and init_labels
M0 = opts.init_cmeans;
if ~isempty(M0)
    assert(isnumeric(M0) && isequal(size(M0), [d K]), ...
        'sltoolbox:slkmeans:invalidopt', ...
        'The option init_cmeans should be a d x K numeric matrix.');
end
L0 = opts.init_labels;
if ~isempty(L0)
    assert(isnumeric(L0) && isvector(L0) && length(L0) == n, ...
        'sltoolbox:slkmeans:invalidopt', ...
        'The option init_labels should be a number vector of length n.');
    assert(all(L0 >= 1) && all(L0 == fix(L0)) && all(L0 <= K), ...
        'sltoolbox:slkmeans:invalidopt', ...
        'The values in option init_labels should be an integer in [1, K]');
    
    if ~isempty(M0)
        error('sltoolbox:slkmeans:invalidopt', ...
            'Cannot simultaneously specify both the options init_cmeans and init_labels.');
    end    
end


% weights
weights = opts.weights;
if ~isempty(weights)
    assert(isnumeric(weights) && isvector(weights) && length(weights) == n, ...
        'sltoolbox:slkmeans:invalidopt', ...
        'The option weights should be a numeric vector of length n.');
end

% handle_empty
assert(ischar(opts.handle_empty), ...
    'sltoolbox:slkmeans:invalidopt', 'The option handle_empty should be a string.');
assert(any(strcmp(opts.handle_empty, {'error', 'retain', 'singleton', 'drop'})), ...
    'sltoolbox:slkmeans:invalidopt', 'Unknown name for the handle_empty: %s', opts.handle_empty);
        
% metric and compute
if isempty(opts.compute)    % use metric 
    assert(ischar(opts.metric), ...
        'sltoolbox:slkmeans:invalidopt', 'The option metric should be a string.');
    
    switch opts.metric
        case 'sqdist'
            f_preproc = [];
            f_metric = @(x1, x2) slmetric_pw(x1, x2, 'sqdist');
            f_centroid = @slmean;
        case 'corrdist'
            f_preproc = @slnrmvecs;
            f_metric = @(x1, x2) 1 - x1' * x2;
            f_centroid = @(x, w) slnrmvecs(slmean(x, w));
        case 'cityblk'
            f_preproc = [];
            f_metric = @(x1, x2) slmetric_pw(x1, x2, 'cityblk');
            f_centroid = @(x, w) median(x, 2);
            
            if ~isempty(weights)
                error('sltoolbox:slkmeans:invalidopt', ...
                    'Sample weighting is not supported when using metric cityblk.');
            end
        otherwise
            error('sltoolbox:slkmeans:invalidopt', ...
                'The option metric can not be set to %s', opts.metric);
    end
    
else                        % use compute    
    assert(iscell(opts.compute) && numel(opts.compute) == 2 && ...
        isa(opts.compute{1}, 'function_handle') && ...
        isa(opts.compute{2}, 'function_handle'), ...
        'sltoolbox:slkmeans:invalidopt', ...
        'The option compute should be a cell array with two function handles.');
    
    f_preproc = [];
    f_metric = opts.compute{1};
    f_centroid = opts.compute{2};
end

% start
assert(ischar(opts.start), ...
    'sltoolbox:slkmeans:invalidopt', 'The option start should be a string.');
switch opts.start
    case 'rndsample';   f_start = @(x) x(:, randsample(n, K));
    case 'farsample';   f_start = @(x) x(:, select_farthest(x, f_metric, K));
    otherwise
        error('sltoolbox:slkmeans:invalidopt', ...
            'Unknown name for start option: %s', opts.start);
end


% maxiter
maxiter = opts.maxiter;
assert(isnumeric(maxiter) && isscalar(maxiter) && maxiter > 1 && maxiter == fix(maxiter), ...
    'sltoolbox:slkmeans:invalidopt', 'The option maxiter should be an integer larger than 1.');

% verbose
verbose = opts.verbose;
assert(islogical(verbose) && isscalar(verbose), ...
    'sltoolbox:slkmeans:invalidopt', 'The option verbose should be a logical scalar.');

output_info = (nargout >= 4);


%% main 

% initialization

if verbose
    show_info(0, 'K-means Initialization.');
end

if ~isempty(f_preproc)
    X = f_preproc(X);
end

if ~isempty(L0)
    cmeans = [];
    labels = L0;
    ginds = slgroups(labels, 1:K);
    cellfun(@(u) assert(~isempty(u)), ginds, 'sltoolbox:slkmeans:invalidopt', ...
        'The init_labels should cover 1:K.');
    
elseif ~isempty(M0)
    cmeans = M0;
    mvs = f_metric(cmeans, X);
    [labels, ginds, dists] = assign_labels(mvs, K);
    
else
    cmeans = f_start(X);
    mvs = f_metric(cmeans, X);
    [labels, ginds, dists] = assign_labels(mvs, K);
end

% preparation

if output_info
    changes = zeros(maxiter, 1);
    avgdists = zeros(maxiter, 1);
end

niters = 0;
c_updated = true(1, K);


if verbose
    show_info(0, 'Entering K-means iteration.');
end

while niters < maxiter
    
    niters = niters + 1;
    
    if verbose
        show_info(1, 'Iter %d  (K = %d):', niters, K);
    end
    
    % handle empty clusters
    
    if ~isempty(cmeans)        
        c_isempty = cellfun(@isempty, ginds);
        if any(c_isempty)      % if there are some empty clusters, take proper actions
            switch opts.handle_empty
                case 'error'
                    error('sltoolbox:slkmeans:emptyclus', ...
                        'Empty clusters are not allowed.');
                    
                case 'retain'
                    c_updated(c_isempty) = false;
                    
                case 'singleton'
                    inds_single_samples = choose_singleton_samples(dists, ginds, c_isempty);
                    
                    c_updated(c_isempty) = true;
                    c_updated(labels(inds_single_samples)) = true;
                    
                    labels(inds_single_samples) = find(c_isempty);
                    ginds = slgroups(labels, 1:K);
                    assert(~any(cellfun(@isempty, ginds)));
                    
                    cmeans(:, c_isempty) = X(:, inds_single_samples);                    
                    
                    if verbose
                        show_info(2, 'apply create new singletons for %d empty clusters', ...
                            sum(c_isempty));
                    end
                    
                case 'drop'
                    ne = sum(c_isempty);
                    labelmap = zeros(1, K);
                    labelmap(~c_isempty) = 1:(K-ne);
                    
                    labels = labelmap(labels);
                    cmeans(:, c_isempty) = [];
                    ginds(c_isempty) = [];
                    c_updated(c_isempty) = [];
                    mvs(c_isempty, :) = [];                    
                    K = K - ne;
                    
                    if verbose
                        show_info(2, 'drop %d empty clusters', ...
                            sum(c_isempty));
                    end
            end
        end
    end
    
    % update means
    
    if all(c_updated)
        cmeans = compute_centroids(X, weights, ginds, f_centroid);
    else
        cmeans(:, c_updated) = compute_centroids(X, weights, ginds(c_updated), f_centroid);
    end    
    
    % update labels
    
    prev_labels = labels;    
    if isempty(mvs)
        mvs = f_metric(cmeans, X);
    else
        mvs(c_updated, :) = f_metric(cmeans(:, c_updated), X);
    end        
    [labels, ginds, dists] = assign_labels(mvs, K);
    
    s_updated = (prev_labels ~= labels);
    cur_change = sum(s_updated);
    
    % record
    
    if output_info
        changes(niters) = cur_change;
        if isempty(weights)
            avgdists(niters) = mean(dists);
        else
            avgdists(niters) = sum(dists * weights(:)) / sum(weights);
        end
    end
    
    if verbose
        show_info(2, 'changes = %d,   avg distances = %g', ...
            cur_change, mean(dists));
    end
    
    if cur_change == 0
        break;
    else
        c_updated = false(1, K);
        c_updated(prev_labels(s_updated)) = true;
        c_updated(labels(s_updated)) = true;
    end           
    
end

if verbose
    if cur_change == 0
        show_info(0, 'converged.');
    else
        show_info(0, 'NOT converged.');
    end
end


if output_info
    info = struct( ...
        'nsamples', n, ...
        'nclusters', K, ...
        'initK', opts.K, ...
        'niters', niters, ...
        'converged', cur_change == 0, ...
        'changes', changes(1:niters), ...
        'avgdists', avgdists(1:niters));
end

    

%% Core functions

function [labels, ginds, dists] = assign_labels(mvs, K)

[dists, labels] = min(mvs, [], 1);
ginds = slgroups(labels, 1:K);


function cmeans = compute_centroids(X, w, ginds, f_centroid)

if isempty(w)
    mcell = cellfun(@(u) f_centroid(X(:, u), []), ginds, 'UniformOutput', false);
else
    mcell = cellfun(@(u) f_centroid(X(:, u), w(u)), ginds, 'UniformOutput', false);
end

cmeans = [mcell{:}];


%% Auxiliary function

function inds = choose_singleton_samples(dists, ginds, c_isempty)

inds_reserved = cellfun(@(x) select_reserved_idx(dists, x), ginds(~c_isempty));
is_available = true(1, length(dists));
is_available(inds_reserved) = false;
available_inds = find(is_available);

nc_empty = sum(c_isempty);
if nc_empty == 1
    [mindist, inds] = min(dists(available_inds));
    inds = available_inds(inds);
else
    [sdists, inds] = sort(dists(available_inds), 2, 'descend');
    inds = available_inds(inds(1:nc_empty));
end
    

function idx = select_reserved_idx(dists, inds)

if numel(inds) == 1
    idx = inds;
else
    [mdist, idx] = min(dists(inds));
    idx = inds(idx);
end    


function inds = select_farthest(x, f_metric, K)

dists = f_metric(x, x);
[dummy, idx] = max(sum(dists, 1));

if K == 1
    inds = idx;
else
    inds = zeros(1, K);
    inds(1) = idx;
    
    cdists = dists(idx, :);    
    
    for i = 2 : K
        [dummy, idx] = max(cdists);
        inds(i) = idx;
        cdists = min(cdists, dists(idx, :));
    end
end


function show_info(indentlevel, msg, varargin)

sz = [blanks(indentlevel * 4), sprintf(msg, varargin{:})];
disp(sz);



 